'''upernet_beitlarge_ade20k'''
import os
from .base_cfg import *


# modify dataset config
DATASET_CFG = DATASET_CFG.copy()
DATASET_CFG.update({
    'type': 'inria',
    'rootdir': '/data02/ybn1/Datasets/Building/Inria',
})
# DATASET_CFG['train']['aug_opts'] = [
#     ('Resize', {'output_size': (2560, 640), 'keep_ratio': True, 'scale_range': (0.5, 2.0)}),
#     ('RandomCrop', {'crop_size': (640, 640), 'one_category_max_ratio': 0.75}),
#     ('RandomFlip', {'flip_prob': 0.5}),
#     ('PhotoMetricDistortion', {}),
#     ('Normalize', {'mean': [123.675, 116.28, 103.53], 'std': [58.395, 57.12, 57.375]}),
#     ('ToTensor', {}),
#     ('Padding', {'output_size': (128, 128), 'data_type': 'tensor'}),
# ]
# DATASET_CFG['train']['repeat_times'] = 35000
# DATASET_CFG['test']['aug_opts'] = [
#     ('Resize', {'output_size': (300, 300), 'keep_ratio': True, 'scale_range': None}),
#     ('Normalize', {'mean': [123.675, 116.28, 103.53], 'std': [58.395, 57.12, 57.375]}),
#     ('ToTensor', {}),
# ]
# modify dataloader config
DATALOADER_CFG = DATALOADER_CFG.copy()
DATALOADER_CFG['train'].update({
    'batch_size': 16
})
# modify optimizer config
OPTIMIZER_CFG = OPTIMIZER_CFG.copy()
# modify scheduler config
SCHEDULER_CFG = SCHEDULER_CFG.copy()
SCHEDULER_CFG.update({
    'max_epochs': 60,
})
# modify losses config
LOSSES_CFG = LOSSES_CFG.copy()
# modify segmentor config
SEGMENTOR_CFG = SEGMENTOR_CFG.copy()
SEGMENTOR_CFG.update({
    'num_classes': 2,
    'backbone': {
        'type': 'beit_large_patch16_224_pt22k_ft22k',
        'series': 'beit',
        'pretrained': True,
        'selected_indices': (0, 1, 2, 3),
        'norm_cfg': {'type': 'layernorm', 'eps': 1e-6},
    },
    'head': {
        'feature2pyramid': {
            'embed_dim': 1024, 
            'rescales': [4, 2, 1, 0.5],
        },
        'in_channels_list': [1024, 1024, 1024, 1024],
        'feats_channels': 1024,
        'pool_scales': [1, 2, 3, 6],
        'dropout': 0.1,
    },
    'auxiliary': {
        'in_channels': 1024,
        'out_channels': 512,
        'dropout': 0.1,
    }
})
# modify inference config
INFERENCE_CFG = INFERENCE_CFG.copy()
# modify common config
COMMON_CFG = COMMON_CFG.copy()
COMMON_CFG['work_dir'] = 'upernet_beitlarge_Inria_small'
COMMON_CFG['logfilepath'] = 'upernet_beitlarge_Inria_small/upernet_beitlarge_Inria_small.log'
COMMON_CFG['resultsavepath'] = 'upernet_beitlarge_Inria_small/upernet_beitlarge_Inria_small_results.pkl'
COMMON_CFG['save_interval_epochs'] = 10